# Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from websockets.sync.server import serve
import websockets
import struct
from packagetool import PackageTool
from runtimemanager import RuntimeManager
from envmanager import EnvManager
from deploy_context import DeployContext
import os
import platform
import json
import threading
import rospy
import time
import sys


send_locker = threading.Lock()

enableCompress = False
lastWS = None

def do_send(ws, pt):
    global lastWS
    if not lastWS:
        return
    ws = lastWS
    try:
        if enableCompress:
            pt.compress()
        global send_locker
        with send_locker:
            ws.send(struct.pack("<i", len(pt.buffer)))
            ws.send(pt.buffer)
    except Exception as e:
        print("send error:", e)
        ws.close()
        lastWS = None


def proc_echo(ws, pt):
    echoid = pt.get_uint32()

    pt_ret = PackageTool()
    pt_ret.put_string("echo_ret")
    pt_ret.put_uint32(echoid + 1)
    do_send(ws, pt_ret)

    if RuntimeManager.gi().isStateChanged():
        print("update state")
        pt_ret = PackageTool()
        pt_ret.put_string("get_ros_state_ret")
        pt_ret.put_string("ok")
        pt_ret.put_uint32(len(RuntimeManager.gi().nodes))
        for nn in RuntimeManager.gi().nodes.keys():
            pt_ret.put_string(nn)
            pt_ret.put_string_array(RuntimeManager.gi().nodes[nn]["pub"])
            pt_ret.put_string_array(RuntimeManager.gi().nodes[nn]["sub"])
            pt_ret.put_string_array(RuntimeManager.gi().nodes[nn]["srv"])
        do_send(ws, pt_ret)

    RuntimeManager.gi().SendBackMsg()


def proc_get_env(ws, pt):
    global enableCompress
    enableCompress = pt.get_bool()

    ret = platform.uname()
    pt_ret = PackageTool()
    pt_ret.put_string("get_env_ret")
    pt_ret.put_string(
        """system: %s
node name: %s
release: %s
version: %s
machine: %s"""
        % (ret.system, ret.node, ret.release, ret.version, ret.machine)
    )
    pt_ret.put_bool(os.path.exists("/opt/ros/noetic"))
    pt_ret.put_bool(os.path.exists("/opt/ros/rolling"))
    pt_ret.put_bool(os.path.exists("/opt/ros/dora"))
    pt_ret.put_string(os.path.expanduser("~"))
    do_send(ws, pt_ret)


def proc_deploy_context(ws, pt):
    ros_type = pt.get_string()
    cancel = pt.get_bool()

    if cancel:
        DeployContext.gi().start(ros_type, lambda pt_ret: do_send(ws, pt_ret))
    else:
        DeployContext.gi().stop(ros_type, lambda pt_ret: do_send(ws, pt_ret))


def proc_get_project_list(ws, pt):
    workspace_path = pt.get_string()

    if not os.path.exists(workspace_path):
        os.makedirs(workspace_path)

    prj_list = os.listdir(workspace_path)
    pt_ret = PackageTool()
    pt_ret.put_string("get_project_list_ret")
    pt_ret.put_string_array(prj_list)
    do_send(ws, pt_ret)


def proc_create_project(ws, pt):
    workspace_path = pt.get_string()
    project_name = pt.get_string()
    project_detail = pt.get_string()
    project_path = os.path.join(workspace_path, project_name)

    pt_ret = PackageTool()
    pt_ret.put_string("create_project_ret")
    if os.path.exists(project_path):
        pt_ret.put_bool(False)
        pt_ret.put_string("project already exists")
        do_send(ws, pt_ret)
        return
    os.makedirs(project_path)
    project_file = os.path.join(project_path, "maestro_project.json")
    with open(project_file, "w") as f:
        f.write(project_detail)
    pt_ret.put_bool(True)
    pt_ret.put_string(project_path)
    do_send(ws, pt_ret)


def proc_get_file(ws, pt):
    file_path = pt.get_string()

    pt_ret = PackageTool()
    pt_ret.put_string("get_file_ret")
    pt_ret.put_string(file_path)
    if not os.path.exists(file_path):
        print("get_file failed : file not exists", file_path)
        pt_ret.put_bool(False)
        pt_ret.put_string("file not exists")
        do_send(ws, pt_ret)
        return
    with open(file_path, "rb") as f:
        content = f.read()
        pt_ret.put_bool(True)
        pt_ret.put_uint32(len(content))
        pt_ret.put_byte_array(content)
        do_send(ws, pt_ret)


def proc_put_file(ws, pt):
    file_path = pt.get_string()
    size = pt.get_uint32()
    content = pt.get_byte_array()

    pt_ret = PackageTool()
    pt_ret.put_string("put_file_ret")
    prefix_dir, name = os.path.split(file_path)
    if os.path.exists(prefix_dir):
        if not os.path.isdir(prefix_dir):
            print("put_file failed : not a directory", prefix_dir)
            pt_ret.put_bool(False)
            pt_ret.put_string("not a directory")
            do_send(ws, pt_ret)
            return
        elif os.path.isdir(file_path):
            print("put_file failed : file already exists", file_path)
            pt_ret.put_bool(False)
            pt_ret.put_string("file already exists")
            do_send(ws, pt_ret)
            return
    else:
        os.makedirs(prefix_dir)
    with open(file_path, "wb") as f:
        # print("write", content)
        f.write(bytearray(content))
    pt_ret.put_bool(True)
    pt_ret.put_string(file_path)
    do_send(ws, pt_ret)


def proc_get_ros_state(ws, pt):
    RuntimeManager.gi().Expand()
    pt_ret = PackageTool()
    pt_ret.put_string("get_ros_state_ret")
    pt_ret.put_string("ok")
    pt_ret.put_uint32(len(RuntimeManager.gi().nodes))
    for nn in RuntimeManager.gi().nodes.keys():
        pt_ret.put_string(nn)
        pt_ret.put_string_array(RuntimeManager.gi().nodes[nn]["pub"])
        pt_ret.put_string_array(RuntimeManager.gi().nodes[nn]["sub"])
        pt_ret.put_string_array(RuntimeManager.gi().nodes[nn]["srv"])
    do_send(ws, pt_ret)


def proc_query_topic_msg_type(ws, pt):
    topic_name = pt.get_string()

    msgType = RuntimeManager.gi().QueryTopicMsgType(topic_name)
    # print("query_topic_msg_type", topic_name, "ret", msgType)
    pt_ret = PackageTool()
    pt_ret.put_string("query_topic_msg_type_ret")
    pt_ret.put_string("ok")
    pt_ret.put_string(topic_name)
    pt_ret.put_string(msgType)
    do_send(ws, pt_ret)


def proc_query_msg_define(ws, pt):
    msg_type = pt.get_string()
    # print("query_msg_define", msg_type)
    msgObj = EnvManager.gi().GetMsgObj(msg_type)
    pt_ret = PackageTool()
    pt_ret.put_string("query_msg_define_ret")
    pt_ret.put_string(msg_type)
    if msgObj is None:
        pt_ret.put_string("没有找到该消息类型定义:" + msg_type)
    else:
        pt_ret.put_string("ok")
        msgObj.put_struct(pt_ret)
    do_send(ws, pt_ret)

def proc_collect_messages(ws, pt):
    pth = pt.get_string()
    
    pt_ret = PackageTool()
    pt_ret.put_string("collect_messages_ret")
    pt_ret.put_string(pth)
    collect_msgs = []
    for root, dirs, files in os.walk(pth):
        for name in files:
            if name.endswith(".msg") and root.endswith("/msg"):
                collect_msgs.append(root.split("/")[-2]+"/"+name[:-4])
    pt_ret.put_string_array(collect_msgs)
    do_send(ws, pt_ret)
                    

def subscribe_topic_callback(topic_name, msg_type, msg, ws):
    pt = PackageTool()
    pt.put_string("subscribe_topic_ret")
    pt.put_string(topic_name)
    pt.put_string(msg_type)
    msgObj = EnvManager.gi().GetMsgObj(msg_type)
    msgObj.put_msg(pt, msg)
    try:
        # t1 = time.time()
        # print("send 1", pt.offset)
        do_send(ws, pt)
        # print(time.time())
        # print("send 2", time.time() - t1)
    except:
        print("send error")
        return False
    return True


def proc_subscribe_topic(ws, pt):
    topic_name = pt.get_string()
    operate = pt.get_uint8()

    RuntimeManager.gi().SubscribeTopic(
        topic_name, operate, subscribe_topic_callback, ws
    )


publish_threads = {}


def stop_publish_topic(topic_name):
    # print("stop publish_topic", topic_name)
    with send_locker:
        if topic_name in publish_threads:
            ps = publish_threads[topic_name]
            ps["locker"].release()
            ps["publisher"].unregister()
            del ps["publisher"]
            # t=publish_threads[topic_name]["thread"]
            del publish_threads[topic_name]
            # t.join()


def publish_thread_func(topic_name):
    # rospy.init_node('maestro')
    while True:
        with send_locker:
            if topic_name not in publish_threads:
                print("publish_thread_func exit")
                return
            ps = publish_threads[topic_name]
        ps["publisher"].publish(ps["msg"])
        # print("发布了",ps["msg"])
        if ps["hz"] == 0:  # 只发布一次
            print("publish_thread_func once", time.time())
            ps["locker"].acquire()
            # rospy.sleep(0.1)
            # print("publish_thread_func once -")
            # stop_publish_topic(topic_name)
            # return
        else:
            time.sleep(1 / ps["hz"])


def proc_publish_topic(ws, pt):
    topic_name = pt.get_string()
    operate = pt.get_uint8()
    if operate == 0:  # 停止发布
        stop_publish_topic(topic_name)
    elif operate == 1:
        msg_type = RuntimeManager.gi().QueryTopicMsgType(topic_name)
        print("publish_topic", topic_name, msg_type)
        msgObj = EnvManager.gi().GetMsgObj(msg_type)

        nodeName, msgName = msg_type.split("/")
        vars_dict = {
            "emsg": None,
        }
        exec(
            """from %s.msg import %s
emsg = %s()"""
            % (nodeName, msgName, msgName),
            globals(),
            vars_dict,
        )
        # print("  1",msg_type)
        # 填充vars_dict["emsg"]
        try:
            msgObj.get_msg(pt, vars_dict["emsg"])
            hz = pt.get_float32()
            # print("get msg ok", vars_dict["emsg"])
        except Exception as e:
            print("get_msg error", e)
            exit()
        # print("  2",vars_dict["emsg"])

        with send_locker:
            if topic_name in publish_threads:  # 已经启动了发布线程，更新参数
                # print("  4 update msg")
                publish_threads[topic_name]["ws"] = ws
                publish_threads[topic_name]["msg_type"] = msg_type
                publish_threads[topic_name]["msg"] = vars_dict["emsg"]
                publish_threads[topic_name]["hz"] = hz
                publish_threads[topic_name]["locker"].release()
                return

        try:
            exec(
                """
import rospy
from %s.msg import %s
publisher = rospy.Publisher("%s", %s, queue_size=10)
time.sleep(0.3)
"""
                % (nodeName, msgName, topic_name, msgName),
                globals(),
                vars_dict,
            )
        except Exception as e:
            print("publish error", e)
            exit()
        with send_locker:
            if topic_name not in publish_threads:
                # print("  3 start publish_thread")
                publish_threads[topic_name] = {
                    "thread": threading.Thread(
                        target=publish_thread_func, args=(topic_name,)
                    ),
                    "ws": ws,
                    "msg_type": msg_type,
                    "msg": vars_dict["emsg"],
                    "hz": hz,
                    "publisher": vars_dict["publisher"],
                    "locker": threading.Lock(),
                }
                publish_threads[topic_name]["locker"].acquire()
                publish_threads[topic_name]["thread"].start()

from custom_node_manager import CustomNodeManager

def proc_custom_node_operate(ws, pt):
    name = pt.get_string()
    token = pt.get_string()
    operate = pt.get_string()
    fn = pt.get_string()
    code = pt.get_string()
    CustomNodeManager.GetInstance().operate(name, token,operate,fn,code,lambda pt_ret: do_send(ws, pt_ret))


def handler(websocket):
    bytes_data = b""
    afterGetEnv = False  # get_env不压缩
    while True:
        try:
            message = websocket.recv()
        except websockets.exceptions.ConnectionClosed:
            print("connection closed")
            break
        except Exception as e:
            print("recv error", e)
            break

        bytes_data += message

        msg_len = struct.unpack("<i", bytes_data[:4])[0]
        # print("msg_len", msg_len)
        if msg_len <= len(bytes_data) - 4:
            pt = PackageTool(bytes_data[4 : 4 + msg_len])
            if afterGetEnv and enableCompress:
                pt.decompress()
            proto = pt.get_string()
            # if proto != "echo" and proto != "subscribe_topic":
            # print(proto)
            try:
                global lastWS
                if lastWS != websocket:
                    lastWS = websocket
                if proto == "echo":
                    proc_echo(websocket, pt)
                elif proto == "get_env":
                    proc_get_env(websocket, pt)
                    afterGetEnv = True
                elif proto == "deploy_context":
                    proc_deploy_context(websocket, pt)
                elif proto == "get_project_list":
                    proc_get_project_list(websocket, pt)
                elif proto == "create_project":
                    proc_create_project(websocket, pt)
                elif proto == "get_file":
                    proc_get_file(websocket, pt)
                elif proto == "put_file":
                    proc_put_file(websocket, pt)
                elif proto == "get_ros_state":
                    proc_get_ros_state(websocket, pt)
                elif proto == "query_topic_msg_type":
                    proc_query_topic_msg_type(websocket, pt)
                elif proto == "query_msg_define":
                    proc_query_msg_define(websocket, pt)
                elif proto == "collect_messages":
                    proc_collect_messages(websocket, pt)
                elif proto == "subscribe_topic":
                    proc_subscribe_topic(websocket, pt)
                elif proto == "publish_topic":
                    proc_publish_topic(websocket, pt)
                elif proto == "custom_node_operate":
                    proc_custom_node_operate(websocket, pt)
                else:
                    print("unknown proto", proto)
            except:
                print("proc error", proto)
            bytes_data = bytes_data[4 + msg_len :]


# 创建并启动 WebSocket 服务器
def start_server():
    port = 54321
    if len(sys.argv) > 1:
        port = int(sys.argv[1])
    print("start server at 0.0.0.0:%d" % port)
    serve(handler, "0.0.0.0", port).serve_forever()


threading.Thread(target=start_server).start()

input()

print("exit")
os.system(
    """ps -ef | grep "python mock" |grep -v grep | awk '{print $2}' | xargs kill -9"""
)
